//! An extensible MySQL Proxy Server based on tokio-core


#[macro_use]
extern crate log;
extern crate env_logger;
#[macro_use]
extern crate futures;
#[macro_use]
extern crate tokio_core;
extern crate byteorder;

#[macro_use]
extern crate num_derive;

use std::rc::Rc;
use std::io::{self, Read, Write, Error, ErrorKind};
use std::net::Shutdown;

use futures::{Future, Poll, Async};
use tokio_core::net::{TcpStream};
use crate::packet::Packet;
use crate::handle::PacketHandler;
use crate::util::parse_packet_length;

pub mod util;
pub mod packet;
pub mod handle;
pub mod connection;
pub mod channel;

/// Handlers return a variant of this enum to indicate how the proxy should handle the packet.
#[derive(Debug,PartialEq)]
pub enum Action {
    /// drop the packet
    Drop,
    /// forward the packet unmodified
    Forward,
    /// forward a mutated packet
    Mutate(Packet),
    /// respond to the packet without forwarding
    Respond(Vec<Packet>),
    /// respond with an error packet
    Error { code: u16, state: [u8; 5], msg: String },
}








/// Wrapper for TcpStream with some built-in buffering
struct ConnReader {
    stream: Rc<TcpStream>,
    packet_buf: Vec<u8>,
    read_buf: Vec<u8>,
}

/// Wrapper for TcpStream with some built-in buffering
struct ConnWriter {
    stream: Rc<TcpStream>,
    write_buf: Vec<u8>,
}

impl ConnReader {

    fn new(stream: Rc<TcpStream>) -> Self {
        ConnReader {
            stream: stream,
            packet_buf: Vec::with_capacity(4096),
            read_buf: vec![0_u8; 4096]
        }
    }

    /// Read from the socket until the status is NotReady
    fn read(&mut self) -> Poll<(), io::Error> {
        debug!("read()");
        loop {
            match self.stream.poll_read() {
                Async::Ready(_) => {
                    let n = try_nb!((&*self.stream).read(&mut self.read_buf[..]));
                    if n == 0 {
                        return Err(Error::new(ErrorKind::Other, "connection closed"));
                    }
                    self.packet_buf.extend_from_slice(&self.read_buf[0..n]);
                },
                _ => return Ok(Async::NotReady),
            }
        }
    }

    fn next(&mut self) -> Option<Packet> {
        debug!("next()");
        // do we have a header
        if self.packet_buf.len() > 3 {
            let l = parse_packet_length(&self.packet_buf);
            // do we have the whole packet?
            let s = 4 + l;
            if self.packet_buf.len() >= s {
                let p = Packet::new(self.packet_buf.drain(0..s).collect());
                //let p= Packet{ bytes: self.packet_buf.drain(0..s).collect() };
                Some(p)
            } else {
                None
            }
        } else {
            None
        }
    }
}

impl ConnWriter {

    fn new(stream: Rc<TcpStream>) -> Self {
        ConnWriter{
            stream: stream,
            write_buf: Vec::with_capacity(4096),
        }
    }

    /// Write a packet to the write buffer
    fn push(&mut self, p: &Packet) {
        //        debug!("push() capacity: {} position: {} packet_size: {}",
        //               self.write_buf.capacity(), self.write_pos, p.bytes.len());

        self.write_buf.extend_from_slice(&p.bytes);
        debug!("end push()");
    }

    /// Writes the contents of the write buffer to the socket
    fn write(&mut self) -> Poll<(), io::Error> {
        debug!("write()");
        while self.write_buf.len() > 0 {
            match self.stream.poll_write() {
                Async::Ready(_) => {
                    let s = (&*self.stream).write(&self.write_buf[..])?;
                    let _ : Vec<u8> = self.write_buf.drain(0..s).collect();
                },
                _ => return Ok(Async::NotReady)
            }
        }
        return Ok(Async::Ready(()));
    }
}

pub struct Pipe<H: PacketHandler + 'static> {
    client_reader: ConnReader,
    client_writer: ConnWriter,
    server_reader: ConnReader,
    server_writer: ConnWriter,
    handler: H,
}

impl<H> Pipe<H> where H: PacketHandler + 'static {
    pub fn new(client: Rc<TcpStream>,
               server: Rc<TcpStream>,
               handler: H
    ) -> Pipe<H> {

        Pipe {
            client_reader: ConnReader::new(client.clone()),
            client_writer: ConnWriter::new(client),
            server_reader: ConnReader::new(server.clone()),
            server_writer: ConnWriter::new(server),
            handler: handler,
        }
    }
}

impl<H> Future for Pipe<H> where H: PacketHandler + 'static {
    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<(), Error> {
        loop {
            let client_read = self.client_reader.read();

            // process buffered requests
            while let Some(request) = self.client_reader.next() {
                match self.handler.handle_request(&request) {
                    Action::Drop => {},
                    Action::Forward => self.server_writer.push(&request),
                    Action::Mutate(ref p2) => self.server_writer.push(p2),
                    Action::Respond(ref v) => {
                        for p in v {
                            self.client_writer.push(&p);
                        }
                    },
                    Action::Error { code, state, msg } => {
                        let error_packet = Packet::error_packet(code, state, msg);
                        self.client_writer.push(&error_packet);
                    }
                };
            }

            // try reading from server
            let server_read = self.server_reader.read();

            // process buffered responses
            while let Some(response) = self.server_reader.next() {
                match self.handler.handle_response(&response) {
                    Action::Drop => {},
                    Action::Forward => self.client_writer.push(&response),
                    Action::Mutate(ref p2) => self.client_writer.push(p2),
                    Action::Respond(ref v) => {
                        for p in v {
                            self.server_writer.push(&p);
                        }
                    },
                    Action::Error { code, state, msg } => {
                        let error_packet = Packet::error_packet(code, state, msg);
                        self.client_writer.push(&error_packet);
                    }
                };
            }

            // perform all of the writes at the end, since the request handlers may have
            // queued packets in either, or both directions

            // try writing to client
            let client_write = self.client_writer.write();

            // if the server connection has closed, close the client connection too
            match &server_read {
                &Err(ref e) => {
                    debug!("Server closed connection: {}", e);
                    match self.client_writer.stream.shutdown(Shutdown::Write) {
                        Ok(_) => {},
                        Err(_) => {}
                    }
                },
                _ => {}
            }

            // try writing to server
            let server_write = self.server_writer.write();

            // if the client connection has closed, close the server connection too
            match &client_read {
                &Err(ref e) => {
                    debug!("Client closed connection: {}", e);
                    match self.server_writer.stream.shutdown(Shutdown::Write) {
                        Ok(_) => {},
                        Err(_) => {}
                    }
                },
                _ => {}
            }

            try_ready!(client_read);
            try_ready!(client_write);
            try_ready!(server_read);
            try_ready!(server_write);
        }

    }

}


